#include "armor_detector.h"
#include <opencv2/opencv.hpp>
#include <iostream>
#include <sys/stat.h>

ArmorDetector::ArmorDetector() {
    enemy_color_ = 1; // 默认红色
    kernel_ = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(3, 3));
    
    // 默认参数
    binary_thresh_ = 150;
    min_lightbar_area_ = 50;
    max_lightbar_ratio_ = 10;
    min_armor_ratio_ = 1.0;
    max_armor_ratio_ = 5.0;
}

bool ArmorDetector::loadParams(const std::string& config_path) {
    cv::FileStorage fs(config_path, cv::FileStorage::READ);
    if (!fs.isOpened()) {
        std::cout << "Failed to open config file: " << config_path << std::endl;
        return false;
    }
    
    fs["binary_thresh"] >> binary_thresh_;
    fs["min_lightbar_area"] >> min_lightbar_area_;
    fs["max_lightbar_ratio"] >> max_lightbar_ratio_;
    fs["min_armor_ratio"] >> min_armor_ratio_;
    fs["max_armor_ratio"] >> max_armor_ratio_;
    
    fs.release();
    // store config path and last modification time for hot reload
    config_path_ = config_path;
    struct stat st;
    if (stat(config_path.c_str(), &st) == 0) last_mtime_ = st.st_mtime;
    return true;
}

void ArmorDetector::setEnemyColor(int color) {
    enemy_color_ = color;
}

void ArmorDetector::preprocess(const cv::Mat& src, cv::Mat& binary) {
    // try reload params if file changed
    tryReloadParams();

    cv::Mat gray, color_channel;
    std::vector<cv::Mat> channels;

    // protect against single-channel or empty input
    if (src.empty()) {
        binary = cv::Mat();
        return;
    }

    if (src.channels() == 3) {
        cv::cvtColor(src, gray, cv::COLOR_BGR2GRAY);
        cv::split(src, channels);
        if (enemy_color_ == 0) { // 蓝色
            color_channel = channels[0] - channels[2]; // B - R
        } else { // 红色
            color_channel = channels[2] - channels[0]; // R - B
        }
    } else {
        // if single channel, just use it
        src.copyTo(color_channel);
    }
    
    cv::threshold(color_channel, binary, binary_thresh_, 255, cv::THRESH_BINARY);
    cv::morphologyEx(binary, binary, cv::MORPH_CLOSE, kernel_);
}

std::vector<LightBar> ArmorDetector::findLightBars(const cv::Mat& binary) {
    std::vector<LightBar> light_bars;
    std::vector<std::vector<cv::Point>> contours;
    
    cv::findContours(binary, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
    
    for (const auto& contour : contours) {
        double area = cv::contourArea(contour);
        if (area < min_lightbar_area_) continue;
        
        cv::RotatedRect rect = cv::minAreaRect(contour);
    float width = std::min(rect.size.width, rect.size.height);
    float length = std::max(rect.size.width, rect.size.height);
    if (width <= 1e-6f) continue; // avoid div by zero or extremely small
    float ratio = length / width;
        
        if (ratio > max_lightbar_ratio_ || ratio < 2.0) continue;
        
        LightBar light_bar;
        light_bar.rect = rect;
        light_bar.length = length;
        light_bar.width = width;
        light_bar.center = rect.center;
        
        // 获取端点
        cv::Point2f vertices[4];
        rect.points(vertices);
        for (int i = 0; i < 4; i++) {
            light_bar.endpoints.push_back(vertices[i]);
        }
        
        light_bars.push_back(light_bar);
    }
    
    return light_bars;
}

std::vector<Armor> ArmorDetector::matchLightBars(const std::vector<LightBar>& light_bars) {
    std::vector<Armor> armors;
    
    if (light_bars.size() < 2) return armors;
    
    for (size_t i = 0; i < light_bars.size(); i++) {
        for (size_t j = i + 1; j < light_bars.size(); j++) {
            const LightBar& left = light_bars[i];
            const LightBar& right = light_bars[j];
            
            // 计算两个灯条之间的距离和角度差
            float distance = cv::norm(left.center - right.center);
            float avg_length = (left.length + right.length) / 2;
            float ratio = distance / avg_length;
            
            if (ratio < min_armor_ratio_ || ratio > max_armor_ratio_) continue;
            
            // 角度差检查
            float angle_diff = std::abs(left.rect.angle - right.rect.angle);
            if (angle_diff > 20.0) continue;
            
            // 创建装甲板 - 使用 Armor 类构造函数
            std::vector<cv::Point2f> points;
            points.insert(points.end(), left.endpoints.begin(), left.endpoints.end());
            points.insert(points.end(), right.endpoints.begin(), right.endpoints.end());
            
            cv::Rect bounding_rect = cv::boundingRect(points);
            cv::Point2f center = (left.center + right.center) / 2;
            
            Armor armor(bounding_rect, center, points, -1); // -1 表示未识别数字
            
            armors.push_back(armor);
        }
    }
    
    return armors;
}

std::vector<Armor> ArmorDetector::detect(const cv::Mat& frame) {
    cv::Mat binary;
    preprocess(frame, binary);
    
    auto light_bars = findLightBars(binary);
    auto armors = matchLightBars(light_bars);
    
    return armors;
}

void ArmorDetector::drawResults(cv::Mat& frame, const std::vector<Armor>& armors) {
    for (const auto& armor : armors) {
        // Draw bounding box and center
        cv::rectangle(frame, armor.getBoundingRect(), cv::Scalar(0, 255, 0), 2);
        cv::circle(frame, armor.getCenter(), 3, cv::Scalar(0, 0, 255), -1);

        // Draw convex hull of detected points
        auto pts = armor.getPoints();
        if (!pts.empty()) {
            std::vector<cv::Point> hull;
            hull.reserve(pts.size());
            for (auto &p : pts) hull.emplace_back(cv::Point(cvRound(p.x), cvRound(p.y)));
            std::vector<cv::Point> ch;
            cv::convexHull(hull, ch);
            if (!ch.empty()) cv::polylines(frame, ch, true, cv::Scalar(0, 200, 255), 2);
        }

        // Draw number label area
        cv::Rect num_region = armor.getBoundingRect();
        num_region.height = static_cast<int>(num_region.height * 0.3);
        num_region.y = armor.getBoundingRect().y + armor.getBoundingRect().height - num_region.height;
        cv::rectangle(frame, num_region, cv::Scalar(255, 255, 0), 1);

        if (armor.getNumber() != -1) {
            cv::putText(frame, std::to_string(armor.getNumber()), 
                       cv::Point(armor.getBoundingRect().x, armor.getBoundingRect().y + armor.getBoundingRect().height + 20),
                       cv::FONT_HERSHEY_SIMPLEX, 0.8, cv::Scalar(255, 255, 0), 2);
        }
    }
}

void ArmorDetector::tryReloadParams() {
    if (config_path_.empty()) return;
    struct stat st;
    if (stat(config_path_.c_str(), &st) == 0) {
        if (st.st_mtime != last_mtime_) {
            loadParams(config_path_);
            std::cout << "Reloaded params from " << config_path_ << std::endl;
        }
    }
}